from .tinyhouse import Pedigree
from .tinyhouse import InputOutput
from .BlockImputation import gibbs_sampler as ImputeBlock
import numpy as np
import argparse

import concurrent.futures
from itertools import repeat


try:
    profile
except:
    def profile(x): 
        return x

def getArgs() :
    parser = argparse.ArgumentParser(description='')
    core_parser = parser.add_argument_group("Core arguments")
    core_parser.add_argument('-out', required=True, type=str, help='The output file prefix.')


    core_parser.add_argument('-map', required=True, type=str, help='Genetic map (three columns, first one should be chromosome).')
    core_parser.add_argument('-qtl', required=True, type=str, help='QTL file with one phenotype per line, and nLoci columns.')
    core_parser.add_argument('-phenotypes', required=True, type=str, help='Phenotypes for each individual, with one individual per line.')

    core_parser.add_argument('-samplesize', default=1000, type=int, help='Number of haplotypes to sample.')
    core_parser.add_argument('-multiple', default=1, type=int, help='Scaling factor.')
    core_parser.add_argument('-haploid', action='store_true', required=False, help='Option to run the algorithm in a haploid mode.')

    core_parser.add_argument('-maxthreads', default=1, type=int, help='number of threads to use when multithreading.')


    InputOutput.addInputFileParser(parser)
    
    return InputOutput.parseArgs("BlockImpute", parser)

def readMap(mapFile):
    chrs = []
    with open(mapFile) as f:
        for line in f:
            parts = line.split(); 
            try:
                chrs.append(int(parts[0]))
            except ValueError:
                pass
    return chrs
def readQtl(file):
    qtl = []
    with open(file) as f:
        for line in f:
            parts = line.split(); 
            qtl.append(np.array(parts, dtype = np.float32))

    return np.array(qtl, dtype = np.float32)

def readPhenotypes(file, pedigree):
    with open(file) as f:
        for line in f:
            parts = line.split(); 
            idx = parts[0]
            parts = parts[1:]
            phenotype = np.array(parts, dtype = np.float32)

            if idx in pedigree.individuals:
                pedigree.individuals[idx].phenotype = phenotype



@profile
def main():
    args = getArgs() 
    pedigree = Pedigree.Pedigree() 
    # Get genotypes
    InputOutput.readInPedigreeFromInputs(pedigree, args, genotypes = True, haps = True, reads = True)

    # Read in genetic map.
    chrMap = readMap(args.map)
    # Read in QTL effects.
    qtl = readQtl(args.qtl)

    # Read in phenotypes
    readPhenotypes(args.phenotypes, pedigree)

    phenotyped_ind = [ind for ind in pedigree if ind.phenotype is not None]

    if args.maxthreads == 1:
        for ind in phenotyped_ind :
            if args.haploid:
                ImputeBlock.imputeIndividualFromPhenotypes_haploid(ind, sire = ind.sire, dam = ind.dam, qtls = qtl, chrMap = chrMap)
            else:
                ImputeBlock.imputeIndividualFromPhenotypes(ind, sire = ind.sire, dam = ind.dam, qtls = qtl, chrMap = chrMap)

    else:
       with concurrent.futures.ThreadPoolExecutor(max_workers=args.maxthreads) as executor:
            executor.map(impute_ind, phenotyped_ind, repeat(qtl), repeat(chrMap), repeat(args))
 
    pedigree.writeDosages(args.out + ".dosages")




def impute_ind(ind, qtl, chrMap, args):
    if args.haploid:
        ImputeBlock.imputeIndividualFromPhenotypes_haploid(ind, sire = ind.sire, dam = ind.dam, qtls = qtl, chrMap = chrMap)
    else:
        ImputeBlock.imputeIndividualFromPhenotypes(ind, sire = ind.sire, dam = ind.dam, qtls = qtl, chrMap = chrMap)

if __name__ == "__main__":
    main()